from PySide6.QtWidgets import QMainWindow, QWidget, QSpinBox, QSizePolicy, QPushButton,\
                                QLabel, QSpacerItem, QFrame, QHBoxLayout, QVBoxLayout, QColorDialog
from PySide6.QtGui import QPaintEvent, QMouseEvent, QDragEnterEvent, QDropEvent, \
                            QPixmap, QPainter, QColor, QPen, QWheelEvent
from PySide6.QtCore import Qt, QSize, QPoint, QPointF, QRectF, QRect



class UiLayout:

    def __init__(self, parent: QWidget):
        central_widget = QWidget(parent)
        widget = QWidget(central_widget)
        layout = QHBoxLayout(widget)

        lbl1 = QLabel("偏移", widget)
        lbl2 = QLabel("步长:", widget)
        lbl3 = QLabel("区域:", widget)
        lbl4 = QLabel("缩放:", widget)
        lbl5 = QLabel("背景:", widget)

        frame1 = self.create_v_separator(widget)
        frame2 = self.create_v_separator(widget)
        frame3 = self.create_v_separator(widget)
        frame4 = self.create_v_separator(widget)

        spacer = QSpacerItem(40, 20, QSizePolicy.Expanding, QSizePolicy.Minimum)

        self.offset_x_sp = QSpinBox(widget)
        self.offset_y_sp = QSpinBox(widget)
        self.step_x_sp = QSpinBox(widget)
        self.step_y_sp = QSpinBox(widget)
        self.region_lbl = QLabel("(0, 0, 0, 0)", widget)
        self.scaled_lbl = QLabel("x1.0", widget)
        self.bg_btn = QPushButton(widget)

        layout.addWidget(lbl1)
        layout.addWidget(self.offset_x_sp)
        layout.addWidget(self.offset_y_sp)
        layout.addWidget(frame1)
        layout.addWidget(lbl2)
        layout.addWidget(self.step_x_sp)
        layout.addWidget(self.step_y_sp)
        layout.addWidget(frame3)
        layout.addWidget(lbl4)
        layout.addWidget(self.scaled_lbl)
        layout.addWidget(frame4)
        layout.addWidget(lbl5)
        layout.addWidget(self.bg_btn)
        layout.addWidget(frame2)
        layout.addWidget(lbl3)
        layout.addWidget(self.region_lbl)
        layout.addItem(spacer)

        widget.setFixedHeight(45)
        
        layout_2 = QVBoxLayout(central_widget)
        layout_2.setContentsMargins(0, 0, 0, 0)
        spacer_2 = QSpacerItem(20, 40, QSizePolicy.Minimum, QSizePolicy.Expanding)

        layout_2.addWidget(widget)
        layout_2.addItem(spacer_2)

        self.offset_x_sp.setPrefix("x")
        self.offset_y_sp.setPrefix("y")
        self.step_x_sp.setPrefix("x")
        self.step_y_sp.setPrefix("y")

        self.step_x_sp.setValue(16)
        self.step_y_sp.setValue(16)
        self.bg_btn.setStyleSheet("background-color: rgb(47, 47, 47)")

        font = self.region_lbl.font()
        font.setPixelSize(24)
        font.setBold(True)
        self.region_lbl.setFont(font)

        widget.setObjectName("head")
        widget.setStyleSheet("#head{background-color: rgba(47, 47, 47, 200)}")

        parent.setCentralWidget(central_widget)

    def create_v_separator(self, parent: QWidget) -> QFrame:
        frame = QFrame(parent)
        frame.setFrameShape(QFrame.Shape.VLine)
        frame.setFrameShadow(QFrame.Shadow.Sunken)
        return frame


class MainWindow(QMainWindow):

    class AtlasData:

        def __init__(self, pixmap: QPixmap):
            self.atlas = pixmap
            self.pos = QPoint()

    
    class Settings:

        def __init__(self):
            self.border_width = 2
            self.border_color = QColor(255, 0, 0)
            self.grid_line_width = 0.5
            self.grid_line_color = QColor(255, 255, 255, 80)
            self.selected_gl_width = 2
            self.selected_gl_color = QColor(0, 255, 0)
            self.scale_delta = 0.1


    class DragAtlasData:

        def __init__(self):
            self.can_drag: bool = False
            self.start_mouse_pos: QPoint = QPoint()
            self.start_atlas_pos: QPoint = QPoint()
    

    class SelectRegionData:

        def __init__(self):
            self.can_select: bool = False
            self.region: QRect = QRect()


    def __init__(self):
        super().__init__()

        self.ui = UiLayout(self)
        self.atlas = self.AtlasData(None)
        self.settings = self.Settings()
        self.drag_atlas = self.DragAtlasData()
        self.select_region = self.SelectRegionData()
        self.scale = 1.0
        self.bg = QColor(47, 47, 47)

        self.setWindowTitle("Atlas Region Selector")
        self.resize(960, 640)
        self.setAcceptDrops(True)

        self.ui.bg_btn.clicked.connect(self.on_pick_bg_color)

    def dragEnterEvent(self, event: QDragEnterEvent):
        if event.mimeData().hasUrls():
            event.acceptProposedAction()

    def dropEvent(self, event: QDropEvent):
        path = event.mimeData().urls()[0].toLocalFile()
        self.atlas.atlas = QPixmap(path)
        # 设置图片位置，使其中心点位于坐标系原点
        self.atlas.pos.setX(-self.atlas.atlas.width() / 2)
        self.atlas.pos.setY(-self.atlas.atlas.height() / 2)
    
    def wheelEvent(self, event: QWheelEvent):
        if self.atlas.atlas is None or self.atlas.atlas.isNull(): return
        delta = event.angleDelta().y()
        
        if delta > 0:
            self.scale += self.settings.scale_delta
        else:
            self.scale -= self.settings.scale_delta

        self.scale = 0.5 if self.scale < 0.5 else 10.0 if self.scale > 10.0 else self.scale
        self.ui.scaled_lbl.setText("x%.1f" % self.scale)

        return super().wheelEvent(event)

    def mousePressEvent(self, event: QMouseEvent):
        if event.button() == Qt.MouseButton.MiddleButton:
            self.drag_atlas.can_drag = True
            self.drag_atlas.start_mouse_pos = event.position().toPoint()
            self.drag_atlas.start_atlas_pos = self.atlas.pos
        
        elif event.button() == Qt.MouseButton.LeftButton and \
            self.has_mouse_point(QRect(self.atlas.pos, self.atlas.atlas.size()), event.position()):
            self.select_region.can_select = True
            # 获取鼠标基于图片左上角为原点的位置
            pos = self.map_mouse_pos(event.position().toPoint()) / self.scale - self.atlas.pos
            self.select_region.region.setX(pos.x())
            self.select_region.region.setY(pos.y())
            self.select_region.region.setSize(QSize(0, 0))

    def mouseReleaseEvent(self, event: QMouseEvent):
        if event.button() == Qt.MouseButton.MiddleButton:
            self.drag_atlas.can_drag = False
        
        elif event.button() == Qt.MouseButton.LeftButton:
            self.select_region.can_select = False
            self.adjust_selected_region()

    def mouseMoveEvent(self, event: QMouseEvent):
        if self.drag_atlas.can_drag:
            offset = event.position() - self.drag_atlas.start_mouse_pos
            offset.setX(offset.x() / self.scale)
            offset.setY(offset.y() / self.scale)
            self.atlas.pos = self.drag_atlas.start_atlas_pos + offset.toPoint()
        
        if self.select_region.can_select:
            offset = self.map_mouse_pos(event.position().toPoint()) / self.scale - self.select_region.region.topLeft() - self.atlas.pos
            self.select_region.region.setWidth(offset.x())
            self.select_region.region.setHeight(offset.y())


    def paintEvent(self, event: QPaintEvent):
        painter = QPainter(self)
        painter.fillRect(self.rect(), self.bg)
        if self.atlas.atlas is None or self.atlas.atlas.isNull(): return
        
        painter.setRenderHints(QPainter.RenderHint.Antialiasing | QPainter.RenderHint.LosslessImageRendering)
        self.keep_center_scale(painter)
        self.draw_atlas(painter)
        self.draw_grid(painter)
        self.draw_selected_region(painter)

    def keep_center_scale(self, painter: QPainter) -> None:
        center = self.rect().center()
        # 将窗口中心位置设为绘制的坐标系的原点
        painter.translate(center.x(), center.y())
        painter.scale(self.scale, self.scale)

    def draw_atlas(self, painter: QPainter) -> None:
        painter.drawPixmap(self.atlas.pos, self.atlas.atlas)
        self.update()

    def draw_grid(self, painter: QPainter) -> None:
        # 绘制边框
        pen = QPen(self.settings.border_color)
        pen.setWidth(self.settings.border_width)
        painter.setPen(pen)
        painter.drawRect(self.atlas.pos.x() - self.settings.border_width,
                         self.atlas.pos.y() - self.settings.border_width, 
                         self.atlas.atlas.width() + self.settings.border_width * 2,
                         self.atlas.atlas.height() + self.settings.border_width * 2)

        # 绘制网格
        pen.setColor(self.settings.grid_line_color)
        pen.setWidth(self.settings.grid_line_width)
        painter.setPen(pen)
        for x in range(self.ui.offset_x_sp.value(), self.atlas.atlas.width(), self.ui.step_x_sp.value()):
            x += self.atlas.pos.x()
            painter.drawLine(x, self.atlas.pos.y(), x, self.atlas.atlas.height() + self.atlas.pos.y())

        for y in range(self.ui.offset_y_sp.value(), self.atlas.atlas.height(), self.ui.step_y_sp.value()):
            y += self.atlas.pos.y()
            painter.drawLine(self.atlas.pos.x(), y, self.atlas.atlas.width() + self.atlas.pos.x(), y)

    def draw_selected_region(self, painter: QPainter) -> None:
        pen = QPen(self.settings.selected_gl_color)
        pen.setWidth(self.settings.selected_gl_width)
        painter.setPen(pen)
        painter.drawRect(self.select_region.region.x() + self.atlas.pos.x(),
                         self.select_region.region.y() + self.atlas.pos.y(),
                         self.select_region.region.width(),
                         self.select_region.region.height())
    
    # 将鼠标的坐标系原点转换为窗口的中心点
    def map_mouse_pos(self, pos: QPoint) -> QPointF:
        """将鼠标的坐标系原点转换为窗口的中心点"""
        return pos - self.rect().center()
    
    def has_mouse_point(self, rect: QRect, pos: QPointF) -> bool:
        p = self.map_mouse_pos(pos)
        p.setX(p.x() / self.scale)
        p.setY(p.y() / self.scale)
        return rect.contains(p.toPoint())

    def adjust_selected_region(self) -> None:
        x = round(self.select_region.region.x() / self.ui.step_x_sp.value()) * self.ui.step_x_sp.value() + self.ui.offset_x_sp.value()
        y = round(self.select_region.region.y() / self.ui.step_y_sp.value()) * self.ui.step_y_sp.value() + self.ui.offset_y_sp.value()
        w = round(self.select_region.region.width() / self.ui.step_x_sp.value()) * self.ui.step_x_sp.value()
        h = round(self.select_region.region.height() / self.ui.step_y_sp.value()) * self.ui.step_y_sp.value()
        self.select_region.region.setTopLeft(QPoint(x, y))
        self.select_region.region.setSize(QSize(w, h))

        self.ui.region_lbl.setText(f"({x}, {y}, {w}, {h})")

    def on_pick_bg_color(self) -> None:
        color = QColorDialog.getColor()
        if color.isValid():
            self.ui.bg_btn.setStyleSheet(f"background-color: {color.name()}")
            self.bg = color



if __name__ == "__main__":
    from PySide6.QtWidgets import QApplication

    app = QApplication()
    win = MainWindow()
    win.show()
    app.exec()